# prototypical network
# Prototypical Networks for Few-shot Learning
# http://arxiv.org/abs/1703.05175

import torch
from torch import nn
from lib.base.base_net import BaseNet


# batch_size: 10
# episode_size: 2000
# batch_num: 2000/10=200
# train_n_way: 60
# train_n_shot: 5
# train_n_query: 5
# val_n_way: 5
# val_n_shot: 5
# val_n_query: 15

class ProtoNet(BaseNet):
    def __init__(self, cfg):
        super().__init__(cfg)

        self.flatten = nn.Flatten()

    # protonet 作者代码 https://github.com/jakesnell/prototypical-networks
    def forward(self, s_q_set, type):
        batch_size = s_q_set.shape[0]  # 10
        n_samples = s_q_set.shape[1]  # 600

        # s_q_set: 10*600*1*28*28 -> 6000*1*28*28
        s_q_set = s_q_set.flatten(start_dim=0, end_dim=1)

        # feature: 6000*dim
        feature = self.flatten(self.backbone(s_q_set))
        # feature: 10*600*dim
        feature = feature.view(batch_size, n_samples, -1)

        # support_sets_f: 10*300*dim
        support_sets_f = feature[:, :self.n_way[type] * self.n_shot[type], :]
        # support_sets_f: 10*60*5*dim
        support_sets_f = support_sets_f.view(batch_size, self.n_way[type], self.n_shot[type], -1)
        # support_sets_f: 10*60*dim
        support_sets_f = torch.mean(support_sets_f, dim=2)
        # query_sets_f: 10*300*dim
        query_sets_f = feature[:, self.n_way[type] * self.n_shot[type]:, :]

        data_f = torch.cat([support_sets_f, query_sets_f], dim=1)

        return data_f
